# coding: utf-8
# Author：WangTianRui
# Date ：2022/3/5 15:53
import Model_ResNetSE34.ResNetSE34V2 as ResNetSE34V2
import Model_ResNetSE34.ResNetSE34L as ResNetSE34L
import torch, os
import utils.features as features_utils
import utils.utils_ as utils
import copy, pickle
import numpy as np
from tqdm import tqdm


def get_wavs():
    all_wavs = utils.get_all_wavs(r"G:\datas\test\SR")
    all_infos = {}
    for i, item in enumerate(all_wavs):
        if item.split("\\")[-2] not in all_infos.keys():
            all_infos[item.split("\\")[-2]] = [item]
        else:
            all_infos[item.split("\\")[-2]].append(item)
    return all_infos


def get_embedding(model, wav_path):
    sig = features_utils.librosa.load(wav_path, sr=16000)[0]
    embedding = model(torch.tensor([sig]))
    return embedding[0].detach().numpy()


if __name__ == '__main__':
    # embedder_pt = torch.load(r"G:\models\SR_models\baseline_v2_ap.model", map_location="cpu")
    embedder_pt = torch.load(r"G:\models\SR_models\baseline_lite_ap.model", map_location="cpu")
    # embedder = ResNetSE34V2.ResNetSE()
    embedder = ResNetSE34L.ResNetSE()
    # print(embedder_pt)
    for key in list(embedder_pt.keys()):
        if str(key).startswith('__S__'):
            embedder_pt[key.replace("__S__.", "")] = embedder_pt[key]
            embedder_pt.pop(key, '404')
        if str(key).startswith('__L__'):
            embedder_pt.pop(key, '404')
    # print(embedder_pt)
    embedder.load_state_dict(embedder_pt)
    embedder.eval()
    info_save_path = "baseline_lite_ap.log"
    if os.path.exists(info_save_path):
        with open(info_save_path, 'rb') as f:
            speaker_with_embedding = pickle.load(f)
    else:
        all_infos = get_wavs()
        speaker_with_embedding = copy.copy(all_infos)
        for reader in all_infos.keys():
            for i, wav_path in enumerate(all_infos[reader][:7]):
                temp = get_embedding(embedder, wav_path)
                speaker_with_embedding[reader][i] = temp
            speaker_with_embedding[reader] = speaker_with_embedding[reader][:7]
        with open(info_save_path, 'wb') as f:
            pickle.dump(speaker_with_embedding, f)
    infos = {}
    print(speaker_with_embedding.keys())
    for speaker in speaker_with_embedding.keys():
        one_embeddings = np.concatenate([x.reshape(1, -1) for x in speaker_with_embedding[speaker]], axis=0).T
        utils.plot_mesh(one_embeddings, title=speaker)
        infos[speaker] = {
            "center": np.mean(one_embeddings, axis=1),
            "all_utts": one_embeddings
        }

    # 开始计算mixture的
    wavs = utils.get_all_wavs(r"G:\datas\test\SR\mixture")
    reference_center = infos["reader_00927"]["center"]
    snrs = []
    neibors = []
    diffs = []
    xs = []
    for wav_path in tqdm(wavs):
        snr = wav_path.split("_snr")[1].split(".")[0]
        neibor = wav_path.split("_snr")[0].split("+")[1]
        inp_embedding = get_embedding(embedder, wav_path)
        diff = np.mean((inp_embedding - reference_center) ** 2)

        snrs.append(snr)
        neibors.append(neibor)
        diffs.append(diff)
        xs.append(neibor[7:] + "_snr_" + str(snr))
    snrs = np.array(snrs)
    neibors = np.array(neibors)
    diffs = np.array(diffs)
    xs = np.array(xs)

    sorted_index = np.argsort(diffs)
    sorted_neibors = neibors[sorted_index]
    sorted_diffs = diffs[sorted_index]
    sorted_snrs = snrs[sorted_index]
    sorted_xs = xs[sorted_index]

    print(sorted_neibors)
    print(sorted_diffs)
    print(sorted_snrs)
    utils.plt.bar(sorted_xs, height=sorted_diffs)
    utils.plt.show()
